// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/plan/plan_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>

using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::functions::aggregate::test;

namespace kumo::pollux::aggregate::test {

class CountIfAggregationTest : public AggregationTestBase {
 protected:
  void SetUp() override {
    AggregationTestBase::SetUp();
  }

  RowTypePtr rowType_{
      ROW({"c0", "c1", "c2"}, {INTEGER(), BOOLEAN(), BOOLEAN()})};
};

TEST_F(CountIfAggregationTest, countIfConst) {
  auto vectors = {
      make_row_vector({
          make_flat_vector<int64_t>(
              10, [](vector_size_t row) { return row / 3; }),
          make_constant(true, 10),
          make_constant(false, 10),
      }),
  };

  createDuckDbTable(vectors);

  testAggregations(
      vectors,
      {},
      {"count_if(c1)", "count_if(c2)"},
      "SELECT sum(if(c1, 1, 0)), sum(if(c2, 1, 0)) FROM tmp");

  testAggregations(
      vectors,
      {"c0"},
      {"count_if(c1)", "count_if(c2)"},
      "SELECT c0, sum(if(c1, 1, 0)), sum(if(c2, 1, 0)) FROM tmp group by c0");
}

TEST_F(CountIfAggregationTest, oneAggregateSingleGroup) {
  // Make two batches of rows: one with nulls; another without.
  auto vectors = {
      make_row_vector(
          {make_flat_vector<bool>(1'000, [](auto row) { return row % 5 == 0; })}),
      make_row_vector({make_flat_vector<bool>(
          1'100, [](auto row) { return row % 3 == 0; }, null_every(7))}),
  };

  createDuckDbTable(vectors);

  testAggregations(
      vectors, {}, {"count_if(c0)"}, "SELECT sum(if(c0, 1, 0)) FROM tmp");
}

TEST_F(CountIfAggregationTest, oneAggregateMultipleGroups) {
  auto vectors = makeVectors(rowType_, 10, 100);
  createDuckDbTable(vectors);

  testAggregations(
      vectors,
      {"c0"},
      {"count_if(c1)"},
      "SELECT c0, sum(if(c1, 1, 0)) FROM tmp GROUP BY c0");
}

TEST_F(CountIfAggregationTest, twoAggregatesSingleGroup) {
  auto vectors = makeVectors(rowType_, 10, 100);
  createDuckDbTable(vectors);

  testAggregations(
      vectors,
      {},
      {"count_if(c1)", "count_if(c2)"},
      "SELECT sum(if(c1, 1, 0)), sum(if(c2, 1, 0)) FROM tmp");
}

TEST_F(CountIfAggregationTest, twoAggregatesMultipleGroups) {
  auto vectors = makeVectors(rowType_, 10, 100);
  createDuckDbTable(vectors);

  testAggregations(
      vectors,
      {"c0"},
      {"count_if(c1)", "count_if(c2)"},
      "SELECT c0, SUM(if(c1, 1, 0)), SUM(if(c2, 1, 0)) FROM tmp GROUP BY c0");
}

TEST_F(CountIfAggregationTest, twoAggregatesMultipleGroupsWrapped) {
  auto vectors = makeVectors(rowType_, 10, 100);
  createDuckDbTable(vectors);

  testAggregations(
      [&](auto& builder) {
        builder.values(vectors)
            .filter("c0 % 2 = 0")
            .project({"c0 % 11 AS c0_mod_11", "c1", "c2"});
      },
      {"c0_mod_11"},
      {"count_if(c1)", "count_if(c2)"},
      "SELECT c0 % 11, SUM(if(c1, 1, 0)), SUM(if(c2, 1, 0)) FROM tmp WHERE c0 % 2 = 0 GROUP BY 1");
}

} // namespace kumo::pollux::aggregate::test
